Human Summary

The host side of HBG was taking 20x more time than the device side. This was due to an unnecessary resetting (zeroing) and copying of the entire scheduling workspace. This PR reduces the H2D transfer to the minimal required information, and the work structures are initialized-on-use, rather than pre-zeroed.

• Up to 99% (depending on the case) of the host-side costs are reduced.
• This optimization is orthogonal to that of Add Graph Execution to host_build_graph #1444, both will help independently in ameliorating the host-side overhead.

AI Summary

PR: host_build_graph — init-on-write SM + runtime arena, bounded to total_tasks

Branch hbg-sm-init-on-write off upstream/main (b7141748). One commit, +121/−34
across 6 files.

TL;DR

HBG's per-dispatch wall is 96–99.6% host bind, and bind was dominated by
rebuilding and H2D-uploading full, ring-sized host structures every run — the
shared-memory mirror and the ~20 MB prebuilt runtime arena — even though a run touches
a tiny fraction of either. Their sizes track ring capacity (task window, 65536-slot pools),
not the workload. The device boots scheduler-only and reads no slot past total_tasks, so
both are made init-on-write and shipped bounded to the task count.

Result — host bind, A/B vs upstream (a2a3, exclusive card, 100-round median):

Device time is unchanged (~0.05–0.1 ms throughout) — the entire win drops into
per-dispatch latency. bind has ~±1 ms run-to-run jitter, but against 77–403 ms baselines
the reductions are unambiguous.

Problem

Splitting bind (via [STRACE] markers) put nearly all of it in run_host_orchestration /
bind_callable_to_runtime_impl, in two structures that are both sized by ring capacity and
rebuilt+reuploaded every run:

• The shared-memory mirror (host_sm_buf(sm_size, 0)): 81 MB (default ring) to 651 MB
  (4 GB ring), of which ~97% is the payload segment.
• The prebuilt runtime arena (~20 MB, fixed): scheduler ready queues (65536×7) +
  tensormap (65536 entries).

A run uses a few dozen of the 16k–131k slots, so almost all of the alloc/zero/init and the
uploads is work on capacity that is never touched.

Fix

The device reads no slot past total_tasks. So:

Shared memory — descriptors and payloads are written per task at submit; slot_states and
completion_flags are reset per slot in orch::prepare_task as it is claimed (dropping the
window-wide reset loop in init_header_per_ring); only the header is zeroed on the host;
each segment is H2D-uploaded bounded to [0, total_tasks).

Runtime arena —

• Don't upload the orchestrator block (fanin_seen_epoch / scope / tensormap, ~8.5 MB): it is
  host-only dep-computation scratch the AICPU scheduler never reads.
• Skip the build's O(capacity) ready-queue slot init (headers only) and seed each big queue's
  slots post-orchestration, uploading only that live prefix instead of the ring-sized tail.
  Safe: relocate_host_orch_image does not walk the ready queues (roots ride in the SM; the
  device boot scan classifies), so uninitialized slots are never read in between.
• Prefix length is min(total_tasks + 1, capacity). A queue takes at most total_tasks
  pushes, but PTO2ReadyQueue::pop_batch_tagged reads one slot past dequeue_pos — the slot
  at enqueue_pos — to detect the empty boundary; a batched dequeue that finds a stale
  (too-large) Vyukov sequence there spins forever. So the prefix seeds one sentinel slot past
  the task count, giving the boundary read a valid "empty" sequence; the tail beyond it is
  never touched. (Capacity is a fixed 65536 while the task window can exceed it, so the prefix
  is clamped.)
• Drop the redundant per-entry stores in the tensormap reset (the preceding memset already
  zeroes the link pointers and producer_task_id).

total_tasks is range-checked before it sizes the copies. Single-ring (PTO2_MAX_RING_DEPTH == 1).

The +1 sentinel slot — why bounding a Vyukov queue to exactly total_tasks deadlocks

This is the non-obvious footgun in bounding the ready queues, worth calling out for anyone
who touches this code or applies the same "ship a live prefix" pattern to another MPMC queue.

A first cut seeded and shipped each ready queue's slots [0, total_tasks) — the reasoning
being "a queue takes at most total_tasks pushes, and reads no slot past its push count." The
first half is true; the second is not. PTO2ReadyQueue::pop_batch_tagged sizes a batch by
scanning forward from dequeue_pos until it hits a slot whose Vyukov sequence marks it empty
— and that terminating read lands on the slot at enqueue_pos, i.e. one slot past the last
one actually pushed. When a queue holds every task, enqueue_pos reaches total_tasks, so
the boundary read touches slot total_tasks — the first slot the [0, total_tasks) prefix did
not seed or upload.

On a fresh device that slot happens to read as empty and nothing breaks — which is why this
passed single-shot runs. On reuse of the persistent device arena (2nd+ invocation of any
callable), that slot still holds a stale, too-large sequence from a prior run. pop_batch
reads it, computes diff > 0, and — per its contract — treats that as "a concurrent producer
is mid-push here, retry" and spins the outer loop forever. The last ready task is never
dequeued; the scheduler latches SCHEDULER_TIMEOUT (surfaced host-side as 507018). It
reproduced intermittently under pytest-xdist (~2 hangs/pass) and was invisible to golden runs
that don't reuse across differently-sized callables.

Fix: seed and ship the prefix as min(total_tasks + 1, capacity) — one sentinel slot past
the task count, carrying its normal empty sequence so the boundary read resolves to "empty" and
pop_batch breaks cleanly. pop_batch never scans past that first empty slot, and per-queue
enqueue_pos ≤ total_tasks, so a single sentinel is sufficient for every queue and every split
of tasks across shapes. The tail beyond it is genuinely never read.

Takeaway for future bounding work: a lock-free queue's read set can extend one element past
its write set. Bound to the read set (pushes + 1), not the write set (pushes).

Why it is safe

Every scheduler-read field is explicitly initialized at submit, so nothing depends on the
removed blanket zeros; reads are bounded by current_task_index. The orchestrator block is
host-only (verified: zero AICPU-scheduler references); relocation never walks the ready
queues; the tensormap-reset change is a pure dedup of what memset already wrote.

Validation

100-round golden × 4 workloads PASS (bgemm, matmul, vector, paged_attention — the last with
its 131072-slot window, exercising the capacity clamp and cross-run queue reuse).

Cross-run reuse is stress-tested under pytest-xdist (one worker per device, the whole HBG
suite, multiple back-to-back passes across 6 devices): 0 scheduler stalls / 0 op-execute
timeouts. This is the regression barrier for the pop_batch boundary slot — before the
total_tasks + 1 sentinel, the same stress produced an intermittent SCHEDULER_TIMEOUT
(≈2 hangs/pass) on the second-and-later invocation of a callable that reused the persistent
device arena.

Full HBG a2a3 scene-test suite green except one failure —
run_stream_reuse::test_depth_two_slots_own_separate_resources (slot-1 GM heap bank
uncommitted) — which reproduces identically on clean upstream/main, i.e. pre-existing and
unrelated (arena-bank path, not this change).

Scope / not done

• a2a3 host_build_graph only. Orthogonal to Add Graph Execution to host_build_graph #1444 (graph execution): that shrinks what goes
  into the SM; this fixes the SM/arena buffer lifecycle. They compound.
• Fully bounding the tensormap reset (a further ~1–2 ms) is deliberately left out:
  print_stats scans the whole entry pool and dereferences entries, and the pool has a
  recycling free-list, so bounding its reset is high-risk for a sub-noise gain. The 8 MB
  memset stays.